# pyright: reportInvalidTypeForm=false
from atomic_agents.lib.factories.mcp_tool_factory import fetch_mcp_tools
from rich.console import Console
from rich.table import Table
import openai
import os
import instructor
from pydantic import Field
from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig
from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
from atomic_agents.lib.components.agent_memory import AgentMemory
from typing import Union, Type, Dict, List
from dataclasses import dataclass
from mcp import ClientSession, StdioServerParameters
import google.generativeai as genai
from mcp.client.stdio import stdio_client
from dotenv import load_dotenv
from Agent.AgentTools import *
current_dir = os.path.dirname(__file__)
project_root = os.path.abspath(os.path.join(current_dir, '..'))
env_path = os.path.join(project_root, ".env")
load_dotenv(env_path)
PROXY = os.getenv("PROXY")
if PROXY:
    os.environ["http_proxy"] = PROXY
    os.environ["https_proxy"] = PROXY 
    os.environ["NO_PROXY"] = "127.0.0.1,localhost"
API_KEY = os.getenv("GEMINI_API_KEY")
console = Console()
tools = get_tools()
# Add FinalResponseSchema for final responses
class FinalResponseSchema(BaseIOSchema):
    """Schema for final response to user."""
    response: str = Field(..., description="The comprehensive final response to the user's query.")
# Build mapping from input_schema to ToolClass
tool_schema_to_class_map: Dict[Type[BaseIOSchema], Type[BaseAgent]] = {
    ToolClass.input_schema: ToolClass for ToolClass in tools if hasattr(ToolClass, "input_schema")
}
# Collect all tool input schemas
tool_input_schemas = tuple(tool_schema_to_class_map.keys())
# Define the Union of all action schemas (including final response)
ActionUnion = Union[tool_input_schemas + (FinalResponseSchema,)]
# 2. Schema and class definitions
class JudgeAgentInputSchema(BaseIOSchema):
    """Schema for the input to the Judge Agent."""
    query: str = Field(..., description="The user's query to analyze and respond to.")
class JudgeAgentOutputSchema(BaseIOSchema):
    """Schema for the output of the Judge Agent."""
    reasoning: str = Field(
        ..., description="Detailed explanation of why this action was chosen and how it will address the user's query."
    )
    action: ActionUnion = Field(
        ..., description="The chosen action: either a tool's input schema instance or a final response schema instance."
    )
def validate_tool_params(tool_class, params):
    """验证工具参数的有效性"""
    try:
        tool_class.input_schema(**params)
        return True
    except Exception as e:
        console.print(f"[red]Parameter validation failed:[/red] {str(e)}")
        return False
def execute_tool_chain(orchestrator_agent, initial_query: str, max_iterations: int = 5) -> str:
    """执行工具链，支持多步骤工具调用"""
    current_query = JudgeAgentInputSchema(query=initial_query)
    for iteration in range(max_iterations):
        console.print(f"[dim]Iteration {iteration + 1}/{max_iterations}[/dim]")
        try:
            orchestrator_output = orchestrator_agent.run(current_query)
            action_instance = orchestrator_output.action
            reasoning = orchestrator_output.reasoning
            console.print(f"[cyan]Agent reasoning:[/cyan] {reasoning}")
            # 检查是否是最终响应
            if isinstance(action_instance, FinalResponseSchema):
                console.print(f"[bold Yellow]This is the final response to the user's query.[/bold Yellow] ")
                return action_instance.response
            # 执行工具
            schema_type = type(action_instance)
            ToolClass = tool_schema_to_class_map.get(schema_type)
            if not ToolClass:
                raise ValueError(f"Unknown schema type '{schema_type.__name__}' returned by orchestrator")
            tool_name = ToolClass.__name__
            console.print(f"[blue]Executing tool:[/blue] {tool_name}")
            params = action_instance.model_dump()
            console.print(f"[dim]Parameters:[/dim] {params}")
            # 验证参数
            if not validate_tool_params(ToolClass, params):
                raise ValueError(f"Invalid parameters for tool {tool_name}")
            # 执行工具
            tool_output = execute_tool(tool_name, **params)
            if hasattr(tool_output, 'model_dump'):
                tool_result_dict = tool_output.model_dump()
                tool_result_str = str(tool_result_dict)
            else:
                tool_result_str = str(tool_output)
            console.print(f"[bold green]Tool Result:[/bold green] {tool_result_str}")
            # 添加工具结果到内存
            tool_result_message = f"Tool {tool_name} executed with parameters {params}. Result: {tool_result_str}"
            orchestrator_agent.memory.add_message("system", JudgeAgentInputSchema(query=tool_result_message))
            # 创建下一轮查询
            current_query = JudgeAgentInputSchema(
                query=f"Based on the tool execution result from {tool_name}: {tool_result_str}. "
                      f"Please either use another tool if more data is needed, or provide a comprehensive final response to the original query: {initial_query}"
            )
        except Exception as e:
            console.print(f"[red]Error in iteration {iteration + 1}:[/red] {str(e)}")
            if iteration == max_iterations - 1:
                return f"Error: Failed to complete analysis after {max_iterations} iterations. Last error: {str(e)}"
            continue
    return f"Maximum iterations ({max_iterations}) reached without final response"
def main():
    try:
        console.print("[bold green]Initializing Judge Agent...[/bold green]")
        # Display available tools
        table = Table(title="Available Tools", box=None)
        table.add_column("Tool Name", style="cyan")
        table.add_column("Input Schema", style="yellow")
        for ToolClass in tools:
            schema_name = ToolClass.input_schema.__name__ if hasattr(ToolClass, "input_schema") else "N/A"
            table.add_row(ToolClass.__name__, schema_name)
        console.print(table)
        # Create and initialize Judge Agent
        console.print("[dim]• Creating Judge Agent...[/dim]")
        memory=AgentMemory()
        orchestrator_agent = BaseAgent(
            BaseAgentConfig(
                client=instructor.from_openai(openai.OpenAI(
                    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
                    api_key=API_KEY,
                )),
                memory=memory,
                model="gemini-2.5-flash",
                system_prompt_generator=SystemPromptGenerator(
                background = [
                    "You are a CTF Performance Evaluation Agent specializing in comprehensive assessment of user performance.",
                    "You have access to multiple analysis tools and should use them strategically to provide thorough evaluations.",
                    "Your goal is to provide fair, objective, and actionable feedback based on data-driven analysis.",
                    "You can execute multiple tools in sequence to gather comprehensive data before providing final responses."
                ],
                steps = [
                    "1. Parse the user's evaluation request and identify required analysis dimensions.",
                    "2. Determine the optimal sequence of tools to gather necessary data.",
                    "3. Execute tools in logical order, using results from previous tools to inform subsequent calls.",
                    "4. Once sufficient data is collected, use FinalResponseSchema to provide comprehensive analysis.",
                    "5. Always validate tool parameters before execution and handle errors gracefully."
                ],
                output_instructions = [
                    "1. Always explain your reasoning clearly before taking action.",
                    "2. Use multiple tools when necessary to provide comprehensive analysis.",
                    "3. Return FinalResponseSchema only when you have sufficient data to answer the user's query.",
                    "4. Ensure all tool parameters are valid and complete.",
                    "5. Provide specific, actionable feedback in final responses.",
                    "6. Maintain professional and constructive tone throughout."
                ]
                ),
                input_schema=JudgeAgentInputSchema,
                output_schema=JudgeAgentOutputSchema,
            )
        )
        console.print("[green]Successfully created Judge Agent.[/green]")
        # 显示可用工具的建议
        console.print()
        # Interactive chat loop
        console.print("[bold green]Judge Agent Interactive Chat. Type 'exit' or 'quit' to leave.[/bold green]")
        while True:
            query = console.input("[bold yellow]You:[/bold yellow] ").strip()
            if query.lower() in {"exit", "quit"}:
                console.print("[bold red]Exiting chat. Goodbye![/bold red]")
                break
            if not query:
                continue  # Ignore empty input
            try:
                result = execute_tool_chain(orchestrator_agent, query)
                console.print(f"[bold green]Judge Agent:[/bold green] {result}")
            except Exception as e:
                console.print(f"[red]Error processing query:[/red] {str(e)}")
                console.print_exception()
    except Exception as e:
        console.print(f"[bold red]Fatal error:[/bold red] {str(e)}")
        console.print_exception()
        return
if __name__ == "__main__":
    main()
